#include "zlib.h"
#include "zconf.h"
#include "bzlib.h"
#include "Client.h"
#include "Packet.h"

//Encryption class
void Client::Codec::Reset(uint8_t key) {
	m_key = key;
	m_offset = 0;
	m_iterator = ITERATOR_START_VAL;
	m_limit = -1;
}

void Client::Codec::Apply(uint8_t* buf, size_t len)
{
	const uint8_t* iterator = reinterpret_cast<const uint8_t*>(
		&m_iterator);

	for (size_t i = 0; i < len; ++i) {
		const size_t i_ = i + m_offset;
		if (i_ % 4 == 0) {
			if (m_limit == 0) return;
			m_iterator *= 0x8088405;
			m_iterator += m_key;
			m_offset = 0;
			if (m_limit > 0) m_limit--;
		}

		buf[i] ^= iterator[i_%4];
	}
}

void Client::Codec::Limit(int32_t limit)
{
	m_limit = limit;
}

int Client::Codec::LimitFromType(uint8_t type)
{
	static int limits[][3] = {
	    {0x02, 0x0C}, {0x04, 0x04}, {0x06, 0x04}
    };

	for (unsigned int i = 0; i < 3; ++i)
	{
		// If we found a valid type, set the limit.
		if (limits[i][0] == type)
		{
			m_limit = (int32_t)limits[i][1];
			return 0;
		}
	}

	// Error.
	return 1;
}

ByteStream Client::DecryptPacket(Client::Codec& codec, int codecType, const char* data, size_t dataLen)
{
    ByteStream retVal(data, dataLen);
    //Configure the decrypter
    codec.LimitFromType(codecType);
    //Decrypt the packet
    codec.Apply(reinterpret_cast<uint8_t*>(&retVal[0]), retVal.Length());
    return retVal;
}


//Client class
void Client::WritePacket(const ByteStream& packet)
{
    m_sendStream.Write(&packet[0], packet.Length());
    m_sendStream.Write('\n');

    if(m_sendStream.Length() > 0x1000)
        CompressStream();
}

void Client::CompressStream()
{
    char buffer[0xFFFF];

    if(m_sendStream.Length())
    {
        if(m_sendStreamFinal.Length() < 0x200000)
        {
            bool compressed = false;

            //If this is the old packet method, then just use zlib compression
            if(m_packetVersion == PACKET_VERSIONOLD)
            {
                size_t bufferLen = sizeof(buffer) - 2;
                compressed = (compress((Bytef*)&buffer[2], (uLongf*)&bufferLen, (const Bytef*)&m_sendStream[0], m_sendStream.Length()) == Z_OK);
                if(compressed)
                {
                    buffer[0] = bufferLen >> 8;
                    buffer[1] = bufferLen & 0xFF;
                    m_sendStreamFinal.Write(buffer, bufferLen + 2);
                }
            } else {
                size_t bufferLen = sizeof(buffer) - 3;
                int codecType = 0;

                //If the packet is bigger than 8KB then use bzlib
                if(m_sendStream.Length() > 0x2000)
                {
                    compressed = (BZ2_bzBuffToBuffCompress(&buffer[3], &bufferLen, &m_sendStream[0], m_sendStream.Length(), 1, 0, 30) == BZ_OK);
                    codecType = Codec::ENCRYPT22_BZ2;
                //Use zlib if the packet is bigger than 50Bytes but small than 8KB
                } else if(m_sendStream.Length() > 50)
                {
                    compressed = (compress((Bytef*)&buffer[3], (uLongf*)&bufferLen, (const Bytef*)&m_sendStream[0], m_sendStream.Length()) == Z_OK);
                    codecType = Codec::ENCRYPT22_ZLIB;
                //Otherwise do not compress (compressing will probably just make it bigger)
                } else {
                    bufferLen = (m_sendStream.Length() < bufferLen ? m_sendStream.Length() : bufferLen);
                    codecType = Codec::ENCRYPT22_UNCOMPRESSED;
                    memcpy(&buffer[3], &m_sendStream[0], bufferLen);
                    compressed = true;
                }

                if(compressed)
                {
                    m_outCodec.LimitFromType(codecType);
                    m_outCodec.Apply(reinterpret_cast<uint8_t*>(&buffer[3]), bufferLen);
                    buffer[0] = (bufferLen + 1) >> 8;
                    buffer[1] = (bufferLen + 1) & 0xFF;
                    buffer[2] = codecType;
                    m_sendStreamFinal.Write(buffer, bufferLen + 3);
                }
            }

            if(!compressed)
            {
                //log << LogHeader() << "Compression error\n";
                Disconnect("Compression error");
            }
            m_sendStream.Clear();
        } else {
            //log << LogHeader() << "Too much data in the send stream (> 2 MegaBytes)\n";
            Disconnect();
        }
    }
}

bool Client::SendPackets()
{
    //If there is data that hasn't been compressed yet, then compress it before
    //we try and send the send stream
    if(m_sendStream.Length())
        CompressStream();

    while(m_sendStreamFinal.Length())
    {
        size_t chunkSize = (m_sendStreamFinal.Length() < 0x2000 ? m_sendStreamFinal.Length() : 0x2000);

        if(m_socket->Write(&m_sendStreamFinal[0], chunkSize) > 0)
        {
            if(chunkSize < m_sendStreamFinal.Length())
                m_sendStreamFinal = ByteStream(&m_sendStreamFinal[chunkSize], m_sendStreamFinal.Length() - chunkSize);
            else m_sendStreamFinal.Clear();
        } else return false;
    }
    return true;
}

bool Client::ReadPacket(ByteStream* stream)
{
    char buffer[0xFFFF];

    if(m_recvStream.Length() >= 2)
    {
        //Get the size of the first packet in the receive stream
        unsigned int dataLen = ((unsigned char)m_recvStream[0] << 8) + (unsigned char)m_recvStream[1];

        //Make sure we have the full packet before trying to process it
        //(otherwise we try and process it later on when we receive the full packet)
        if(dataLen + 2 <= m_recvStream.Length())
        {
            size_t cLen = sizeof(buffer);
            bool decompressed = false;

            //If we're using the old packet method, then uncompress packet with zlib
            if(m_packetVersion == PACKET_VERSIONOLD)
                decompressed = uncompress((Bytef*)buffer, (uLongf*)&cLen, (const Bytef*)&m_recvStream[2], dataLen) == Z_OK;
            else {
                //If we're using the new packet method, then get the compression method
                int codecType = m_recvStream[2];

                //Decrypt the packet
                ByteStream packet = DecryptPacket(m_inCodec, codecType, &m_recvStream[3], dataLen - 1);

                //Use ZLIB to decompress
                if(codecType == Codec::ENCRYPT22_ZLIB)
                    decompressed = (uncompress((Bytef*)buffer, (uLongf*)&cLen, (const Bytef*)&packet[0], packet.Length()) == Z_OK);

                //Use BZLIB to decompress
                else if(codecType == Codec::ENCRYPT22_BZ2)
                    decompressed = (BZ2_bzBuffToBuffDecompress(buffer, &cLen, &packet[0], packet.Length(), 0, 0) == BZ_OK);

                //Use no compression
                else if(codecType == Codec::ENCRYPT22_UNCOMPRESSED)
                {
                    //Just copy the packet to buffer
                    cLen = (packet.Length() < sizeof(buffer) ? packet.Length() : sizeof(buffer));
                    memcpy(buffer, &packet[0], cLen);
                    decompressed = true;
                } else {
                    //log << LogHeader() << "Invalid packet\n";
                    Disconnect("Your client sent an invalid packet.");
                }
            }

            if(decompressed)
            {
                m_recvStream = ByteStream(&m_recvStream[dataLen + 2], m_recvStream.Length() - dataLen - 2);
                *stream = ByteStream(buffer, cLen);
                return true;
            } else {
                //log << LogHeader() << "Decompression error\n";
                Disconnect("Decompression error");
            }
        }
    }
    return false;
}

void Client::RecvPackets()
{
    //Make sure the receive stream isn't too full (2 megabytes)
    if(m_recvStream.Length() < 0x200000)
    {
        //Read all the data from the socket
        char buffer[0xFFFF];
        int size = m_socket->Read(buffer, sizeof(buffer));

        //If we received any data from the socket, then append it to the
        //end of the receive stream.
        if(size > 0)
            m_recvStream.Write(buffer, size);

        //Get the last socket error
        int err = soxLastError();
        if(size == 0 || (size < 0 && err != EWOULDBLOCK))
        {
            //log << LogHeader() << "Player left with socket error (" << err << ")\n";
            Disconnect();
        } else {
            m_lastPacketTime = time(NULL);
        }
    } else {
        //log << LogHeader() << "Too much data in the receive stream (> 2 MegaBytes)\n";
        Disconnect("Unable to process your packets");
    }
}

void Client::Disconnect(const std::string& disMessage)
{
    //WritePacket(ByteStream() << PacketInt1(DISMESSAGE) << disMessage);
    Disconnect();
}

void Client::Disconnect() {
    m_deleteMe = true;
}

///GMClient
gmType GMClient::m_userType = GM_NULL;

GMClient::GMClient(gmMachine* machine, sox::TcpSocket* socket):
    Client(socket) {
    m_userObject = machine->AllocUserObject(this, m_userType);
}

GMClient::~GMClient() {
    m_userObject->m_user = NULL;
}

void GMClient::Register(gmMachine* machine)
{
    m_userType = machine->CreateUserType("Client");
}


